import asyncio
import httpx

from motor.resources.instance import Instance, PDRole, Workload, Endpoint, InsStatus, ParallelConfig
from motor.utils.logger import get_logger

logger = get_logger(__name__)

ip = "127.0.0.1"
mgmt_port = 19998
infer_port = 19999

mgmt_client =  httpx.Client(timeout=None, base_url=f"http://{ip}:{mgmt_port}")
infer_client =  httpx.Client(timeout=None, base_url=f"http://{ip}:{infer_port}")

def refresh_instance():
	instances = []
	for i in range(1):
		prefill_id = i
		instance = Instance(
			job_name=f"prefill_instance_{i+1}",
			model_name="test_model",
			id=prefill_id,
			role=PDRole.ROLE_P,
			status=InsStatus.ACTIVE,
			parallel_config=ParallelConfig(dp=2),
			endpoints={
				ip: {
					prefill_id: Endpoint(
						id=j+1,
						ip=ip,
						port=str(8000 + prefill_id + j),
						status="normal"
					) for j in range(1)
				}
			}
		)
		instances.append(instance)

	for i in range(1):
		decode_id = i + 10
		instance = Instance(
			job_name=f"decode_instance_{i+1}",
			model_name="test_model",
			id=decode_id,
			role=PDRole.ROLE_D,
			status=InsStatus.ACTIVE,
			parallel_config=ParallelConfig(dp=2),
			endpoints={
				ip: {
					decode_id: Endpoint(
						id=j+1,
						ip=ip,
						port=str(8000 + decode_id + j),
						status="normal"
					) for j in range(1)
				}
			}
		)
		instances.append(instance)
  
	payload = {
        "event": "add",
		"instances": [inst.model_dump() for inst in instances]
    }

	response = mgmt_client.post(url='/v1/instances/refresh', json=payload)
	logger.info(f"Refresh Response: {response}")

async def chat():
	response = infer_client.post(url='/v1/chat/completions', json={
		"model": "test-model", 
		"messages": [{"role": "user", "content": "Hello from test client!"}],
	})
	logger.info(f"Chat Response: {response}")
	async for line in response.aiter_lines():
		logger.info(f"Received line: {line}")
	logger.info("Finished chat")


if __name__ == "__main__":
	refresh_instance()
	asyncio.run(chat())